import os
import time
from datetime import timedelta
from typing import Dict, Any, Generator, List, Optional
import easyocr
from openai import OpenAI, APIError
from PyPDF2 import PdfReader
from docx import Document
import numpy as np
import cv2
import re


def classify_payslip_type(text: str) -> str:
    """根据工资单内容识别工资单类型"""
    categories = {
        # 按发放周期
        "月度工资单": ["月工资", "月薪", "当月", "本月", "月份工资"],
        "周工资单": ["周工资", "周薪", "本周", "周结"],
        "日工资单": ["日工资", "日薪", "当日", "日结"],
        "年终奖工资单": ["年终奖", "年度奖金", "十三薪", "双薪"],
        "季度奖金单": ["季度奖", "季奖", "季度绩效"],
        
        # 按工作性质
        "正式员工工资单": ["基本工资", "岗位工资", "职务工资", "社保", "公积金"],
        "临时工工资单": ["临时工", "小时工", "计时工资", "临时用工"],
        "实习生工资单": ["实习工资", "实习津贴", "实习补贴"],
        "兼职工资单": ["兼职", "兼职工资", "兼职费"],
        "劳务费单据": ["劳务费", "劳务报酬", "服务费", "咨询费"],
        
        # 按行业特色
        "销售提成工资单": ["提成", "佣金", "销售奖金", "业绩奖"],
        "加班费工资单": ["加班费", "超时费", "节假日加班", "夜班费"],
        "项目奖金单": ["项目奖金", "项目提成", "完工奖"],
        "绩效工资单": ["绩效工资", "绩效奖金", "考核奖金"],
        
        # 特殊类型
        "补发工资单": ["补发", "补缴", "调薪补差", "工资调整"],
        "离职结算单": ["离职", "结算", "最后工作日", "离职补偿"],
    }

    for name, keywords in categories.items():
        if any(keyword in text for keyword in keywords):
            return name
    return "标准工资单"


def extract_payslip_info(text: str) -> Dict[str, Any]:
    """从工资单文本中提取关键信息"""
    info = {
        "员工姓名": None,
        "员工编号": None,
        "部门": None,
        "职位": None,
        "工资期间": None,
        "基本工资": None,
        "实发工资": None,
        "应发工资": None,
        "扣除项目": [],
        "津贴补贴": [],
        "社保公积金": {},
        "个人所得税": None
    }
    
    # 提取员工姓名
    name_patterns = [
        r"姓名[：:]([\S]+)",
        r"员工姓名[：:]([\S]+)",
        r"姓\s*名[：:]\s*([\S]+)"
    ]
    for pattern in name_patterns:
        match = re.search(pattern, text)
        if match:
            info["员工姓名"] = match.group(1).strip()
            break
    
    # 提取员工编号
    id_patterns = [
        r"员工编号[：:]([\S]+)",
        r"工号[：:]([\S]+)",
        r"编号[：:]([\S]+)"
    ]
    for pattern in id_patterns:
        match = re.search(pattern, text)
        if match:
            info["员工编号"] = match.group(1).strip()
            break
    
    # 提取部门
    dept_patterns = [
        r"部门[：:]([\S]+)",
        r"所属部门[：:]([\S]+)",
        r"科室[：:]([\S]+)"
    ]
    for pattern in dept_patterns:
        match = re.search(pattern, text)
        if match:
            info["部门"] = match.group(1).strip()
            break
    
    # 提取工资金额
    salary_patterns = {
        "基本工资": [r"基本工资[：:]\s*([\d,]+\.?\d*)", r"基础工资[：:]\s*([\d,]+\.?\d*)"],
        "实发工资": [r"实发工资[：:]\s*([\d,]+\.?\d*)", r"实际发放[：:]\s*([\d,]+\.?\d*)", r"净工资[：:]\s*([\d,]+\.?\d*)"],
        "应发工资": [r"应发工资[：:]\s*([\d,]+\.?\d*)", r"应发合计[：:]\s*([\d,]+\.?\d*)"],
        "个人所得税": [r"个人所得税[：:]\s*([\d,]+\.?\d*)", r"个税[：:]\s*([\d,]+\.?\d*)"]
    }
    
    for key, patterns in salary_patterns.items():
        for pattern in patterns:
            match = re.search(pattern, text)
            if match:
                amount_str = match.group(1).replace(',', '')
                try:
                    info[key] = float(amount_str)
                except ValueError:
                    pass
                break
    
    return info


class PayslipAnalyzer:
    def __init__(self, api_key: str):
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )
        self._ocr_reader = None

    def _get_ocr_reader(self):
        """延迟初始化OCR读取器"""
        if self._ocr_reader is None:
            self._ocr_reader = easyocr.Reader(['ch_sim', 'en'])
        return self._ocr_reader

    def _extract_from_image(self, file_path: str, reader=None) -> str:
        """从图片文件中提取文本"""
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            # 使用numpy读取文件，支持中文路径
            img_array = np.fromfile(file_path, dtype=np.uint8)
            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            
            if image is None:
                raise ValueError("无法解析图像，cv2.imdecode 失败")
            
            # 使用传入的reader或获取默认reader
            reader = reader or self._get_ocr_reader()
            result = reader.readtext(image, detail=0)
            
            return "\n".join(result).strip()
            
        except Exception as e:
            raise RuntimeError(f"图片OCR识别失败: {str(e)}")

    def _extract_from_pdf(self, file_path: str) -> str:
        """从PDF文件中提取文本"""
        text = ""
        try:
            with open(file_path, 'rb') as file:
                reader = PdfReader(file)
                for page in reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n"
        except Exception as e:
            raise RuntimeError(f"PDF文件读取失败: {str(e)}")
        return text.strip()

    def _extract_from_docx(self, file_path: str) -> str:
        """从DOCX文件中提取文本"""
        try:
            doc = Document(file_path)
            paragraphs = [para.text for para in doc.paragraphs if para.text.strip()]
            return "\n".join(paragraphs)
        except Exception as e:
            raise RuntimeError(f"DOCX文件读取失败: {str(e)}")

    def extract_text(self, file_path: str) -> str:
        """根据文件类型提取文本内容"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == '.pdf':
            return self._extract_from_pdf(file_path)
        elif ext == '.docx':
            return self._extract_from_docx(file_path)
        elif ext in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp']:
            return self._extract_from_image(file_path)
        else:
            raise ValueError(f"不支持的文件格式: {ext}")

    def extract_texts_from_multiple_images(self, image_paths: List[str]) -> str:
        """从多个图片文件中提取文本"""
        reader = self._get_ocr_reader()
        all_text = []
        
        for path in image_paths:
            try:
                text = self._extract_from_image(path, reader)
                if text.strip():
                    all_text.append(text)
            except Exception as e:
                # 记录错误但继续处理其他文件
                print(f"⚠️ 图片处理失败 [{path}]: {str(e)}")
        
        return "\n\n".join(all_text)

    def analyze_payslip_simple(self, text: str, payslip_type: str, extracted_info: Dict[str, Any]) -> str:
        """简化版工资单分析，返回分析结果字符串"""
        if not text.strip():
            raise ValueError("工资单文本内容为空")

        # 构建提取信息的摘要
        info_summary = "\n".join([
            f"- {k}: {v}" for k, v in extracted_info.items() 
            if v is not None and v != [] and v != {}
        ])

        system_prompt = f"""你是一位专业的劳动法律师和薪酬专家，负责分析工资单的合规性和合理性。
该工资单初步识别为：{payslip_type}

已提取的关键信息：
{info_summary}

请简要分析工资单的完整性、规范性和合规性，重点关注：
1. 员工基本信息完整性
2. 工资构成清晰度
3. 扣除项目合理性
4. 计算准确性
5. 合规性评估

请提供简洁的分析结果。"""

        try:
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"请分析以下工资单：\n{text[:10000]}"}
                ],
                temperature=0.3,
                max_tokens=1000
            )
            
            return response.choices[0].message.content.strip()
            
        except APIError as e:
            raise RuntimeError(f"API请求失败: {str(e)}")
        except Exception as e:
            raise RuntimeError(f"分析过程中出错: {str(e)}")

    def analyze_payslip_stream(self, text: str, payslip_type: str, extracted_info: Dict[str, Any]) -> Generator[str, None, Dict[str, Any]]:
        """流式分析工资单（保留原有功能）"""
        if not text.strip():
            raise ValueError("工资单文本内容为空")

        # 构建提取信息的摘要
        info_summary = "\n".join([
            f"- {k}: {v}" for k, v in extracted_info.items() 
            if v is not None and v != [] and v != {}
        ])

        system_prompt = f"""你是一位专业的劳动法律师和薪酬专家，负责分析工资单的合规性和合理性。
该工资单初步识别为：{payslip_type}

已提取的关键信息：
{info_summary}

请严格按照以下要求进行分析：
1. 首先确认工资单类型是否准确
2. 给出总体评价（规范/基本规范/存在问题/严重不规范）
3. 分析工资单的完整性和规范性
4. 检查以下关键项目：
   - 员工基本信息是否完整（姓名、工号、部门等）
   - 工资构成是否清晰（基本工资、津贴、奖金等）
   - 扣除项目是否合理合法（社保、公积金、个税等）
   - 实发金额计算是否正确
   - 工资标准是否符合当地最低工资标准
   - 加班费计算是否符合劳动法规定
   - 社保公积金缴费基数和比例是否合规
   - 个人所得税计算是否正确
5. 指出存在的问题和风险
6. 提供改进建议和维权提示
7. 如发现违法违规情况，说明相关法律依据"""

        try:
            stream = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"请分析以下工资单：\n{text[:15000]}"}
                ],
                temperature=0.3,
                max_tokens=2000,
                stream=True
            )

            collected_content = []
            total_tokens = 0
            
            for chunk in stream:
                if not chunk.choices:
                    continue
                    
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    collected_content.append(delta.content)
                    yield delta.content
                    
                if hasattr(chunk, 'usage') and chunk.usage:
                    total_tokens = chunk.usage.total_tokens

            return {
                "metadata": {
                    "total_tokens": total_tokens, 
                    "complete_response": "".join(collected_content)
                }
            }

        except APIError as e:
            raise RuntimeError(f"API请求失败: {str(e)}")
        except Exception as e:
            raise RuntimeError(f"分析过程中出错: {str(e)}")